#!/usr/bin/python3
# -*- coding: utf-8 -*-

import akshare as ak
import pandas as pd
from datetime import datetime, timedelta


class StockDataCollector:
    def __init__(self):
        pass

    def fetch_fund_flow_rank(self, indicator="今日"):
        """获取个股资金流排名数据"""
        try:
            return ak.stock_individual_fund_flow_rank(indicator=indicator)
        except Exception as e:
            print(f"获取资金流排名数据失败: {e}")
            return pd.DataFrame()

    def fetch_historical_data(self, symbol, start_date, end_date, adjust="qfq"):
        """获取指定股票的历史行情数据"""
        try:
            historical_data = ak.stock_zh_a_daily(symbol=symbol, start_date=start_date, end_date=end_date,
                                                  adjust=adjust)
            if historical_data.empty:
                print(f"未获取到 {symbol} 的历史数据")
            return historical_data
        except Exception as e:
            print(f"获取历史数据失败: {e}")
            return pd.DataFrame()

    def calculate_moving_averages(self, historical_data):
        """计算均线（5日、10日、20日、30日、60日、120日）"""
        if historical_data.empty or 'close' not in historical_data.columns:
            print("历史数据为空或缺少 'close' 列")
            return historical_data

        windows = [5, 10, 20, 30, 60, 120]
        for window in windows:
            historical_data[f'{window}日均线'] = historical_data['close'].rolling(window=window).mean()

        return historical_data

    def fetch_bid_ask_data(self, symbol):
        """获取指定股票的外盘和内盘数据"""
        try:
            stock_bid_ask_df = ak.stock_bid_ask_em(symbol=symbol)

            # 打印获取的盘口数据
            print(f"{symbol} 的盘口数据：\n{stock_bid_ask_df}")

            if not stock_bid_ask_df.empty and {'外盘', '内盘'}.issubset(stock_bid_ask_df['item'].values):
                # 从 DataFrame 中提取外盘和内盘的值
                outer_volume = stock_bid_ask_df.loc[stock_bid_ask_df['item'] == '外盘', 'value'].values[0]
                inner_volume = stock_bid_ask_df.loc[stock_bid_ask_df['item'] == '内盘', 'value'].values[0]
                return pd.Series({'外盘': outer_volume, '内盘': inner_volume})

            print(f"{symbol} 的盘口数据为空或格式不正确")
            return pd.Series()  # 返回一个空的 Series
        except Exception as e:
            print(f"获取 {symbol} 盘口数据失败: {e}")
            return pd.Series()  # 返回一个空的 Series

    def fetch_realtime_stock_data(self):
        """获取所有沪深京 A 股的实时行情数据"""
        try:
            return ak.stock_zh_a_spot_em()
        except Exception as e:
            print(f"获取实时行情失败: {e}")
            return pd.DataFrame()

    def fetch_combined_data(self, symbols, market='sh', indicator="今日", enable_history=True):
        """整合数据收集主方法"""
        fund_flow_rank = self.fetch_fund_flow_rank(indicator)
        realtime_data = self.fetch_realtime_stock_data()

        # 准备历史数据参数
        end_date = datetime.now().strftime("%Y%m%d")
        start_date = (datetime.now() - timedelta(days=220)).strftime("%Y%m%d") if enable_history else None

        result_df = pd.DataFrame()

        for symbol in symbols:
            try:
                symbol_str = str(symbol)

                # 获取实时数据
                realtime_row = realtime_data.loc[realtime_data['代码'] == symbol_str]
                if realtime_row.empty:
                    print(f"跳过 {symbol}：无实时数据")
                    continue

                # 获取资金流数据
                fund_flow_row = fund_flow_rank.loc[fund_flow_rank['代码'] == symbol_str]
                # 打印获取的资金流数据
                print(f"{symbol} 的资金流数据：\n{fund_flow_row}")

                # 删除不需要的列
                columns_to_drop = ['序号', '代码', '名称', '最新价', '今日涨跌幅']
                fund_flow_row = fund_flow_row.drop(columns=columns_to_drop, errors='ignore')

                # 检查如果fund_flow_row为空，则跳过
                if fund_flow_row.empty:
                    print(f"跳过 {symbol}：无资金流数据或数据格式不正确")
                    continue

                # 组合行
                combined_row = pd.concat([realtime_row.iloc[0], fund_flow_row.iloc[0]], axis=0)

                # 添加历史均线数据
                if enable_history:
                    history_symbol = f"{market}{symbol}"
                    historical_data = self.fetch_historical_data(history_symbol, start_date, end_date)

                    if not historical_data.empty:
                        # 检查历史数据长度
                        print(f"{symbol} 的历史数据计算所有均线，仅有 {len(historical_data)} 条记录")
                        historical_data = self.calculate_moving_averages(historical_data)
                        latest_ma = historical_data.iloc[-1][
                                [f'{w}日均线' for w in [5, 10, 20, 30, 60, 120]]
                        ]
                        combined_row = pd.concat([combined_row, latest_ma])

                # 添加盘口数据
                bid_ask = self.fetch_bid_ask_data(symbol)
                # 打印获取的盘口数据
                print(f"{symbol} 的盘口数据：\n{bid_ask}")

                if not bid_ask.empty:
                    combined_row = pd.concat([combined_row, bid_ask])

                result_df = pd.concat([result_df, combined_row.to_frame().T], ignore_index=True)

            except Exception as e:
                print(f"处理 {symbol} 时发生错误: {e}")
                continue

        return result_df


def main():
    collector = StockDataCollector()

    # 上海股票代码
    sample_symbols_1 = ['603688', '600353', '600727']  # 示例股票代码
    combined_data_1 = collector.fetch_combined_data(
        symbols=sample_symbols_1,
        market='sh',
        indicator="今日",
        enable_history=True
    )

    # 深圳股票代码
    sample_symbols_2 = ['002735', '002204', '000422']  # 示例股票代码
    combined_data_2 = collector.fetch_combined_data(
        symbols=sample_symbols_2,
        market='sz',
        indicator="今日",
        enable_history=True
    )

    # 合并两组数据
    combined_data = pd.concat([combined_data_1, combined_data_2], ignore_index=True)

    print("\n整合后的股票数据：")
    print(combined_data)

    # 保存到Excel
    if not combined_data.empty:
        filename = f"stock_data_{datetime.now().strftime('%Y%m%d_%H%M%S')}.xlsx"
        combined_data.to_excel(filename, index=False)
        print(f"\n数据已保存到 {filename}")
    else:
        print("\n未获取到有效数据")


if __name__ == "__main__":
    main()
